"""
21_vector_embeddings.py
=======================
Vector Embeddings for Neo4j Knowledge Graphs

This module covers:
1. Generating embeddings with sentence-transformers
2. OpenAI embeddings
3. Storing embeddings in Neo4j
4. Creating vector indexes
5. Similarity search operations

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional, Union
from dataclasses import dataclass
from abc import ABC, abstractmethod
import hashlib
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver


@dataclass
class EmbeddingConfig:
    """Configuration for embedding generation."""
    model_name: str
    dimensions: int
    batch_size: int = 32
    normalize: bool = True


class EmbeddingProvider(ABC):
    """Abstract base class for embedding providers."""
    
    @property
    @abstractmethod
    def dimensions(self) -> int:
        """Return embedding dimensions."""
        pass
    
    @abstractmethod
    def embed(self, text: str) -> List[float]:
        """Generate embedding for single text."""
        pass
    
    @abstractmethod
    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for multiple texts."""
        pass


class SentenceTransformerEmbeddings(EmbeddingProvider):
    """
    Local embeddings using sentence-transformers.
    
    Popular models:
    - all-MiniLM-L6-v2: 384 dimensions, fast, good quality
    - all-mpnet-base-v2: 768 dimensions, better quality
    - multi-qa-MiniLM-L6-cos-v1: 384 dims, optimized for QA
    
    No API costs, runs locally, good for privacy.
    """
    
    SUPPORTED_MODELS = {
        "all-MiniLM-L6-v2": 384,
        "all-mpnet-base-v2": 768,
        "multi-qa-MiniLM-L6-cos-v1": 384,
        "paraphrase-MiniLM-L6-v2": 384,
        "all-distilroberta-v1": 768,
    }
    
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        """
        Initialize with sentence-transformers model.
        
        Args:
            model_name: Model identifier
        """
        self.model_name = model_name
        self._model = None
        self._dimensions = self.SUPPORTED_MODELS.get(model_name, 384)
    
    @property
    def model(self):
        """Lazy load the model."""
        if self._model is None:
            try:
                from sentence_transformers import SentenceTransformer
                self._model = SentenceTransformer(self.model_name)
            except ImportError:
                raise ImportError(
                    "sentence-transformers required. "
                    "Install: pip install sentence-transformers"
                )
        return self._model
    
    @property
    def dimensions(self) -> int:
        return self._dimensions
    
    def embed(self, text: str) -> List[float]:
        """Generate embedding for text."""
        embedding = self.model.encode(text, normalize_embeddings=True)
        return embedding.tolist()
    
    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for multiple texts."""
        embeddings = self.model.encode(
            texts, 
            normalize_embeddings=True,
            show_progress_bar=len(texts) > 100
        )
        return [e.tolist() for e in embeddings]


class OpenAIEmbeddings(EmbeddingProvider):
    """
    OpenAI embeddings via API.
    
    Models:
    - text-embedding-3-small: 1536 dimensions (default)
    - text-embedding-3-large: 3072 dimensions
    - text-embedding-ada-002: 1536 dimensions (legacy)
    
    Higher quality but requires API key and has costs.
    """
    
    MODELS = {
        "text-embedding-3-small": 1536,
        "text-embedding-3-large": 3072,
        "text-embedding-ada-002": 1536,
    }
    
    def __init__(
        self,
        model_name: str = "text-embedding-3-small",
        api_key: str = None
    ):
        """
        Initialize OpenAI embeddings.
        
        Args:
            model_name: OpenAI embedding model
            api_key: API key (or uses OPENAI_API_KEY env var)
        """
        self.model_name = model_name
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        self._dimensions = self.MODELS.get(model_name, 1536)
        self._client = None
    
    @property
    def client(self):
        """Lazy load OpenAI client."""
        if self._client is None:
            try:
                from openai import OpenAI
                self._client = OpenAI(api_key=self.api_key)
            except ImportError:
                raise ImportError(
                    "openai package required. Install: pip install openai"
                )
        return self._client
    
    @property
    def dimensions(self) -> int:
        return self._dimensions
    
    def embed(self, text: str) -> List[float]:
        """Generate embedding using OpenAI API."""
        response = self.client.embeddings.create(
            model=self.model_name,
            input=text
        )
        return response.data[0].embedding
    
    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for batch."""
        response = self.client.embeddings.create(
            model=self.model_name,
            input=texts
        )
        return [d.embedding for d in response.data]


class VectorStore:
    """
    Manages vector storage and retrieval in Neo4j.
    
    Vector Index Architecture:
    
    ┌─────────────────────────────────────────────────────────────┐
    │                    VECTOR INDEX (HNSW)                      │
    ├─────────────────────────────────────────────────────────────┤
    │                                                             │
    │  Node: Chunk                                                │
    │  Property: embedding                                        │
    │  Dimensions: 384                                            │
    │  Similarity: cosine                                         │
    │                                                             │
    │  ┌─────────────────────────────────────────────────────┐   │
    │  │  HNSW Graph Structure                               │   │
    │  │                                                     │   │
    │  │  Layer 2:  [N1]─────────────────────[N2]           │   │
    │  │             │                        │              │   │
    │  │  Layer 1:  [N1]──[N3]────[N4]──────[N2]           │   │
    │  │             │     │       │         │              │   │
    │  │  Layer 0:  [N1][N5][N3][N6][N4][N7][N2][N8]       │   │
    │  │                                                     │   │
    │  └─────────────────────────────────────────────────────┘   │
    │                                                             │
    │  Search: O(log n) approximate nearest neighbor              │
    │  Memory: O(n * M) where M = max connections                │
    │                                                             │
    └─────────────────────────────────────────────────────────────┘
    """
    
    def __init__(
        self,
        driver=None,
        embedding_provider: EmbeddingProvider = None
    ):
        """
        Initialize vector store.
        
        Args:
            driver: Neo4j driver
            embedding_provider: Provider for generating embeddings
        """
        self.driver = driver or get_driver()
        self.embedding_provider = embedding_provider
    
    # =========================================================
    # INDEX MANAGEMENT
    # =========================================================
    
    def create_vector_index(
        self,
        index_name: str,
        label: str,
        property_name: str,
        dimensions: int,
        similarity_function: str = "cosine"
    ) -> bool:
        """
        Create a vector index.
        
        Args:
            index_name: Unique index name
            label: Node label to index
            property_name: Property containing vectors
            dimensions: Vector dimensions
            similarity_function: 'cosine' or 'euclidean'
            
        Returns:
            True if created successfully
        """
        query = f"""
        CREATE VECTOR INDEX {index_name} IF NOT EXISTS
        FOR (n:{label})
        ON (n.{property_name})
        OPTIONS {{
            indexConfig: {{
                `vector.dimensions`: {dimensions},
                `vector.similarity_function`: '{similarity_function}'
            }}
        }}
        """
        
        with self.driver.session() as session:
            try:
                session.run(query)
                return True
            except Exception as e:
                print(f"Error creating index: {e}")
                return False
    
    def drop_vector_index(self, index_name: str) -> bool:
        """Drop a vector index."""
        query = f"DROP INDEX {index_name} IF EXISTS"
        
        with self.driver.session() as session:
            try:
                session.run(query)
                return True
            except Exception as e:
                print(f"Error dropping index: {e}")
                return False
    
    def list_vector_indexes(self) -> List[Dict[str, Any]]:
        """List all vector indexes."""
        query = """
        SHOW INDEXES
        WHERE type = 'VECTOR'
        YIELD name, labelsOrTypes, properties, state, options
        RETURN name, labelsOrTypes, properties, state, options
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            return [record.data() for record in result]
    
    # =========================================================
    # EMBEDDING STORAGE
    # =========================================================
    
    def store_embedding(
        self,
        label: str,
        node_id: str,
        text: str,
        property_name: str = "embedding"
    ) -> bool:
        """
        Generate and store embedding for a node.
        
        Args:
            label: Node label
            node_id: Node ID
            text: Text to embed
            property_name: Property to store embedding
            
        Returns:
            True if stored
        """
        if self.embedding_provider is None:
            raise ValueError("Embedding provider not configured")
        
        embedding = self.embedding_provider.embed(text)
        
        query = f"""
        MATCH (n:{label} {{id: $id}})
        SET n.{property_name} = $embedding
        RETURN n.id as id
        """
        
        with self.driver.session() as session:
            result = session.run(query, id=node_id, embedding=embedding)
            return result.single() is not None
    
    def store_embeddings_batch(
        self,
        label: str,
        nodes: List[Dict[str, str]],
        text_property: str = "content",
        embedding_property: str = "embedding",
        batch_size: int = 100
    ) -> int:
        """
        Generate and store embeddings for multiple nodes.
        
        Args:
            label: Node label
            nodes: List of {id, content} dictionaries
            text_property: Property containing text
            embedding_property: Property for embedding
            batch_size: Processing batch size
            
        Returns:
            Number of embeddings stored
        """
        if self.embedding_provider is None:
            raise ValueError("Embedding provider not configured")
        
        total_stored = 0
        
        for i in range(0, len(nodes), batch_size):
            batch = nodes[i:i + batch_size]
            texts = [n.get(text_property, "") for n in batch]
            ids = [n.get("id") for n in batch]
            
            # Generate embeddings
            embeddings = self.embedding_provider.embed_batch(texts)
            
            # Store in Neo4j
            query = f"""
            UNWIND $data AS item
            MATCH (n:{label} {{id: item.id}})
            SET n.{embedding_property} = item.embedding
            RETURN count(*) as stored
            """
            
            data = [
                {"id": id_, "embedding": emb}
                for id_, emb in zip(ids, embeddings)
            ]
            
            with self.driver.session() as session:
                result = session.run(query, data=data)
                total_stored += result.single()["stored"]
        
        return total_stored
    
    # =========================================================
    # SIMILARITY SEARCH
    # =========================================================
    
    def similarity_search(
        self,
        index_name: str,
        query_text: str,
        top_k: int = 10,
        min_score: float = 0.0
    ) -> List[Dict[str, Any]]:
        """
        Perform vector similarity search.
        
        Args:
            index_name: Vector index name
            query_text: Text to search for
            top_k: Number of results
            min_score: Minimum similarity score (0-1)
            
        Returns:
            List of similar nodes with scores
        """
        if self.embedding_provider is None:
            raise ValueError("Embedding provider not configured")
        
        # Generate query embedding
        query_embedding = self.embedding_provider.embed(query_text)
        
        return self.similarity_search_by_vector(
            index_name, query_embedding, top_k, min_score
        )
    
    def similarity_search_by_vector(
        self,
        index_name: str,
        query_vector: List[float],
        top_k: int = 10,
        min_score: float = 0.0
    ) -> List[Dict[str, Any]]:
        """
        Search by pre-computed vector.
        
        Args:
            index_name: Vector index name
            query_vector: Query embedding
            top_k: Number of results
            min_score: Minimum similarity
            
        Returns:
            Similar nodes with scores
        """
        query = """
        CALL db.index.vector.queryNodes($index_name, $k, $vector)
        YIELD node, score
        WHERE score >= $min_score
        RETURN 
            node.id AS id,
            node.content AS content,
            score
        ORDER BY score DESC
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                index_name=index_name,
                k=top_k,
                vector=query_vector,
                min_score=min_score
            )
            return [record.data() for record in result]
    
    def hybrid_search(
        self,
        index_name: str,
        query_text: str,
        keyword: str,
        top_k: int = 10,
        vector_weight: float = 0.7
    ) -> List[Dict[str, Any]]:
        """
        Combine vector search with keyword filtering.
        
        Args:
            index_name: Vector index name
            query_text: Text for semantic search
            keyword: Keyword filter
            top_k: Number of results
            vector_weight: Weight for vector vs keyword (0-1)
            
        Returns:
            Hybrid search results
        """
        if self.embedding_provider is None:
            raise ValueError("Embedding provider not configured")
        
        query_embedding = self.embedding_provider.embed(query_text)
        
        query = """
        CALL db.index.vector.queryNodes($index_name, $k * 2, $vector)
        YIELD node, score
        WHERE toLower(node.content) CONTAINS toLower($keyword)
        WITH node, score
        ORDER BY score DESC
        LIMIT $k
        RETURN 
            node.id AS id,
            node.content AS content,
            score
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                index_name=index_name,
                k=top_k,
                vector=query_embedding,
                keyword=keyword
            )
            return [record.data() for record in result]


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Vector Embeddings - Demonstration")
    print("=" * 60)
    
    # Try sentence-transformers first
    print("\n📌 Initializing Embedding Provider")
    print("-" * 40)
    
    try:
        embeddings = SentenceTransformerEmbeddings("all-MiniLM-L6-v2")
        print(f"   Model: all-MiniLM-L6-v2")
        print(f"   Dimensions: {embeddings.dimensions}")
        
        # Test embedding
        test_text = "Diabetes is a chronic metabolic disorder."
        embedding = embeddings.embed(test_text)
        print(f"   Test embedding generated: {len(embedding)} dimensions")
        print(f"   First 5 values: {embedding[:5]}")
        
        # Test batch
        texts = [
            "Treatment options for type 2 diabetes",
            "Insulin resistance and metabolic syndrome",
            "Blood glucose monitoring techniques"
        ]
        batch_embeddings = embeddings.embed_batch(texts)
        print(f"   Batch embeddings: {len(batch_embeddings)} generated")
        
    except ImportError as e:
        print(f"   sentence-transformers not available: {e}")
        embeddings = None
    
    # Vector store demo
    print("\n📌 Vector Store Operations")
    print("-" * 40)
    
    try:
        driver = get_driver()
        store = VectorStore(driver, embeddings)
        
        # List indexes
        indexes = store.list_vector_indexes()
        print(f"   Existing vector indexes: {len(indexes)}")
        for idx in indexes:
            print(f"   - {idx['name']}: {idx['state']}")
        
        # Create index if embeddings available
        if embeddings:
            print("\n   Creating vector index...")
            created = store.create_vector_index(
                index_name="test_vectors",
                label="TestNode",
                property_name="embedding",
                dimensions=embeddings.dimensions,
                similarity_function="cosine"
            )
            print(f"   Index created: {created}")
            
            # Cleanup
            store.drop_vector_index("test_vectors")
            print("   Test index dropped")
        
        driver.close()
        
    except Exception as e:
        print(f"   Error: {e}")
    
    print("\n" + "=" * 60)
    print("✅ Vector embeddings demonstration complete!")
    print("=" * 60)
